{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "7ca21a5a",
   "metadata": {},
   "source": [
    "# Document Classification with UDOP\n",
    "In this notebook, we will give a sequential example on data processing and model IO for a document classification example"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "0bbde2fb",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['<extra_id_0>', '<extra_id_1>', '<extra_id_2>', '<extra_id_3>', '<extra_id_4>', '<extra_id_5>', '<extra_id_6>', '<extra_id_7>', '<extra_id_8>', '<extra_id_9>', '<extra_id_10>', '<extra_id_11>', '<extra_id_12>', '<extra_id_13>', '<extra_id_14>', '<extra_id_15>', '<extra_id_16>', '<extra_id_17>', '<extra_id_18>', '<extra_id_19>', '<extra_id_20>', '<extra_id_21>', '<extra_id_22>', '<extra_id_23>', '<extra_id_24>', '<extra_id_25>', '<extra_id_26>', '<extra_id_27>', '<extra_id_28>', '<extra_id_29>', '<extra_id_30>', '<extra_id_31>', '<extra_id_32>', '<extra_id_33>', '<extra_id_34>', '<extra_id_35>', '<extra_id_36>', '<extra_id_37>', '<extra_id_38>', '<extra_id_39>', '<extra_id_40>', '<extra_id_41>', '<extra_id_42>', '<extra_id_43>', '<extra_id_44>', '<extra_id_45>', '<extra_id_46>', '<extra_id_47>', '<extra_id_48>', '<extra_id_49>', '<extra_id_50>', '<extra_id_51>', '<extra_id_52>', '<extra_id_53>', '<extra_id_54>', '<extra_id_55>', '<extra_id_56>', '<extra_id_57>', '<extra_id_58>', '<extra_id_59>', '<extra_id_60>', '<extra_id_61>', '<extra_id_62>', '<extra_id_63>', '<extra_id_64>', '<extra_id_65>', '<extra_id_66>', '<extra_id_67>', '<extra_id_68>', '<extra_id_69>', '<extra_id_70>', '<extra_id_71>', '<extra_id_72>', '<extra_id_73>', '<extra_id_74>', '<extra_id_75>', '<extra_id_76>', '<extra_id_77>', '<extra_id_78>', '<extra_id_79>', '<extra_id_80>', '<extra_id_81>', '<extra_id_82>', '<extra_id_83>', '<extra_id_84>', '<extra_id_85>', '<extra_id_86>', '<extra_id_87>', '<extra_id_88>', '<extra_id_89>', '<extra_id_90>', '<extra_id_91>', '<extra_id_92>', '<extra_id_93>', '<extra_id_94>', '<extra_id_95>', '<extra_id_96>', '<extra_id_97>', '<extra_id_98>', '<extra_id_99>', '<extra_l_id_0>', '<extra_l_id_1>', '<extra_l_id_2>', '<extra_l_id_3>', '<extra_l_id_4>', '<extra_l_id_5>', '<extra_l_id_6>', '<extra_l_id_7>', '<extra_l_id_8>', '<extra_l_id_9>', '<extra_l_id_10>', '<extra_l_id_11>', '<extra_l_id_12>', '<extra_l_id_13>', '<extra_l_id_14>', '<extra_l_id_15>', '<extra_l_id_16>', '<extra_l_id_17>', '<extra_l_id_18>', '<extra_l_id_19>', '<extra_l_id_20>', '<extra_l_id_21>', '<extra_l_id_22>', '<extra_l_id_23>', '<extra_l_id_24>', '<extra_l_id_25>', '<extra_l_id_26>', '<extra_l_id_27>', '<extra_l_id_28>', '<extra_l_id_29>', '<extra_l_id_30>', '<extra_l_id_31>', '<extra_l_id_32>', '<extra_l_id_33>', '<extra_l_id_34>', '<extra_l_id_35>', '<extra_l_id_36>', '<extra_l_id_37>', '<extra_l_id_38>', '<extra_l_id_39>', '<extra_l_id_40>', '<extra_l_id_41>', '<extra_l_id_42>', '<extra_l_id_43>', '<extra_l_id_44>', '<extra_l_id_45>', '<extra_l_id_46>', '<extra_l_id_47>', '<extra_l_id_48>', '<extra_l_id_49>', '<extra_l_id_50>', '<extra_l_id_51>', '<extra_l_id_52>', '<extra_l_id_53>', '<extra_l_id_54>', '<extra_l_id_55>', '<extra_l_id_56>', '<extra_l_id_57>', '<extra_l_id_58>', '<extra_l_id_59>', '<extra_l_id_60>', '<extra_l_id_61>', '<extra_l_id_62>', '<extra_l_id_63>', '<extra_l_id_64>', '<extra_l_id_65>', '<extra_l_id_66>', '<extra_l_id_67>', '<extra_l_id_68>', '<extra_l_id_69>', '<extra_l_id_70>', '<extra_l_id_71>', '<extra_l_id_72>', '<extra_l_id_73>', '<extra_l_id_74>', '<extra_l_id_75>', '<extra_l_id_76>', '<extra_l_id_77>', '<extra_l_id_78>', '<extra_l_id_79>', '<extra_l_id_80>', '<extra_l_id_81>', '<extra_l_id_82>', '<extra_l_id_83>', '<extra_l_id_84>', '<extra_l_id_85>', '<extra_l_id_86>', '<extra_l_id_87>', '<extra_l_id_88>', '<extra_l_id_89>', '<extra_l_id_90>', '<extra_l_id_91>', '<extra_l_id_92>', '<extra_l_id_93>', '<extra_l_id_94>', '<extra_l_id_95>', '<extra_l_id_96>', '<extra_l_id_97>', '<extra_l_id_98>', '<extra_l_id_99>', '</extra_l_id_0>', '</extra_l_id_1>', '</extra_l_id_2>', '</extra_l_id_3>', '</extra_l_id_4>', '</extra_l_id_5>', '</extra_l_id_6>', '</extra_l_id_7>', '</extra_l_id_8>', '</extra_l_id_9>', '</extra_l_id_10>', '</extra_l_id_11>', '</extra_l_id_12>', '</extra_l_id_13>', '</extra_l_id_14>', '</extra_l_id_15>', '</extra_l_id_16>', '</extra_l_id_17>', '</extra_l_id_18>', '</extra_l_id_19>', '</extra_l_id_20>', '</extra_l_id_21>', '</extra_l_id_22>', '</extra_l_id_23>', '</extra_l_id_24>', '</extra_l_id_25>', '</extra_l_id_26>', '</extra_l_id_27>', '</extra_l_id_28>', '</extra_l_id_29>', '</extra_l_id_30>', '</extra_l_id_31>', '</extra_l_id_32>', '</extra_l_id_33>', '</extra_l_id_34>', '</extra_l_id_35>', '</extra_l_id_36>', '</extra_l_id_37>', '</extra_l_id_38>', '</extra_l_id_39>', '</extra_l_id_40>', '</extra_l_id_41>', '</extra_l_id_42>', '</extra_l_id_43>', '</extra_l_id_44>', '</extra_l_id_45>', '</extra_l_id_46>', '</extra_l_id_47>', '</extra_l_id_48>', '</extra_l_id_49>', '</extra_l_id_50>', '</extra_l_id_51>', '</extra_l_id_52>', '</extra_l_id_53>', '</extra_l_id_54>', '</extra_l_id_55>', '</extra_l_id_56>', '</extra_l_id_57>', '</extra_l_id_58>', '</extra_l_id_59>', '</extra_l_id_60>', '</extra_l_id_61>', '</extra_l_id_62>', '</extra_l_id_63>', '</extra_l_id_64>', '</extra_l_id_65>', '</extra_l_id_66>', '</extra_l_id_67>', '</extra_l_id_68>', '</extra_l_id_69>', '</extra_l_id_70>', '</extra_l_id_71>', '</extra_l_id_72>', '</extra_l_id_73>', '</extra_l_id_74>', '</extra_l_id_75>', '</extra_l_id_76>', '</extra_l_id_77>', '</extra_l_id_78>', '</extra_l_id_79>', '</extra_l_id_80>', '</extra_l_id_81>', '</extra_l_id_82>', '</extra_l_id_83>', '</extra_l_id_84>', '</extra_l_id_85>', '</extra_l_id_86>', '</extra_l_id_87>', '</extra_l_id_88>', '</extra_l_id_89>', '</extra_l_id_90>', '</extra_l_id_91>', '</extra_l_id_92>', '</extra_l_id_93>', '</extra_l_id_94>', '</extra_l_id_95>', '</extra_l_id_96>', '</extra_l_id_97>', '</extra_l_id_98>', '</extra_l_id_99>', '<extra_t_id_0>', '<extra_t_id_1>', '<extra_t_id_2>', '<extra_t_id_3>', '<extra_t_id_4>', '<extra_t_id_5>', '<extra_t_id_6>', '<extra_t_id_7>', '<extra_t_id_8>', '<extra_t_id_9>', '<extra_t_id_10>', '<extra_t_id_11>', '<extra_t_id_12>', '<extra_t_id_13>', '<extra_t_id_14>', '<extra_t_id_15>', '<extra_t_id_16>', '<extra_t_id_17>', '<extra_t_id_18>', '<extra_t_id_19>', '<extra_t_id_20>', '<extra_t_id_21>', '<extra_t_id_22>', '<extra_t_id_23>', '<extra_t_id_24>', '<extra_t_id_25>', '<extra_t_id_26>', '<extra_t_id_27>', '<extra_t_id_28>', '<extra_t_id_29>', '<extra_t_id_30>', '<extra_t_id_31>', '<extra_t_id_32>', '<extra_t_id_33>', '<extra_t_id_34>', '<extra_t_id_35>', '<extra_t_id_36>', '<extra_t_id_37>', '<extra_t_id_38>', '<extra_t_id_39>', '<extra_t_id_40>', '<extra_t_id_41>', '<extra_t_id_42>', '<extra_t_id_43>', '<extra_t_id_44>', '<extra_t_id_45>', '<extra_t_id_46>', '<extra_t_id_47>', '<extra_t_id_48>', '<extra_t_id_49>', '<extra_t_id_50>', '<extra_t_id_51>', '<extra_t_id_52>', '<extra_t_id_53>', '<extra_t_id_54>', '<extra_t_id_55>', '<extra_t_id_56>', '<extra_t_id_57>', '<extra_t_id_58>', '<extra_t_id_59>', '<extra_t_id_60>', '<extra_t_id_61>', '<extra_t_id_62>', '<extra_t_id_63>', '<extra_t_id_64>', '<extra_t_id_65>', '<extra_t_id_66>', '<extra_t_id_67>', '<extra_t_id_68>', '<extra_t_id_69>', '<extra_t_id_70>', '<extra_t_id_71>', '<extra_t_id_72>', '<extra_t_id_73>', '<extra_t_id_74>', '<extra_t_id_75>', '<extra_t_id_76>', '<extra_t_id_77>', '<extra_t_id_78>', '<extra_t_id_79>', '<extra_t_id_80>', '<extra_t_id_81>', '<extra_t_id_82>', '<extra_t_id_83>', '<extra_t_id_84>', '<extra_t_id_85>', '<extra_t_id_86>', '<extra_t_id_87>', '<extra_t_id_88>', '<extra_t_id_89>', '<extra_t_id_90>', '<extra_t_id_91>', '<extra_t_id_92>', '<extra_t_id_93>', '<extra_t_id_94>', '<extra_t_id_95>', '<extra_t_id_96>', '<extra_t_id_97>', '<extra_t_id_98>', '<extra_t_id_99>', '</extra_t_id_0>', '</extra_t_id_1>', '</extra_t_id_2>', '</extra_t_id_3>', '</extra_t_id_4>', '</extra_t_id_5>', '</extra_t_id_6>', '</extra_t_id_7>', '</extra_t_id_8>', '</extra_t_id_9>', '</extra_t_id_10>', '</extra_t_id_11>', '</extra_t_id_12>', '</extra_t_id_13>', '</extra_t_id_14>', '</extra_t_id_15>', '</extra_t_id_16>', '</extra_t_id_17>', '</extra_t_id_18>', '</extra_t_id_19>', '</extra_t_id_20>', '</extra_t_id_21>', '</extra_t_id_22>', '</extra_t_id_23>', '</extra_t_id_24>', '</extra_t_id_25>', '</extra_t_id_26>', '</extra_t_id_27>', '</extra_t_id_28>', '</extra_t_id_29>', '</extra_t_id_30>', '</extra_t_id_31>', '</extra_t_id_32>', '</extra_t_id_33>', '</extra_t_id_34>', '</extra_t_id_35>', '</extra_t_id_36>', '</extra_t_id_37>', '</extra_t_id_38>', '</extra_t_id_39>', '</extra_t_id_40>', '</extra_t_id_41>', '</extra_t_id_42>', '</extra_t_id_43>', '</extra_t_id_44>', '</extra_t_id_45>', '</extra_t_id_46>', '</extra_t_id_47>', '</extra_t_id_48>', '</extra_t_id_49>', '</extra_t_id_50>', '</extra_t_id_51>', '</extra_t_id_52>', '</extra_t_id_53>', '</extra_t_id_54>', '</extra_t_id_55>', '</extra_t_id_56>', '</extra_t_id_57>', '</extra_t_id_58>', '</extra_t_id_59>', '</extra_t_id_60>', '</extra_t_id_61>', '</extra_t_id_62>', '</extra_t_id_63>', '</extra_t_id_64>', '</extra_t_id_65>', '</extra_t_id_66>', '</extra_t_id_67>', '</extra_t_id_68>', '</extra_t_id_69>', '</extra_t_id_70>', '</extra_t_id_71>', '</extra_t_id_72>', '</extra_t_id_73>', '</extra_t_id_74>', '</extra_t_id_75>', '</extra_t_id_76>', '</extra_t_id_77>', '</extra_t_id_78>', '</extra_t_id_79>', '</extra_t_id_80>', '</extra_t_id_81>', '</extra_t_id_82>', '</extra_t_id_83>', '</extra_t_id_84>', '</extra_t_id_85>', '</extra_t_id_86>', '</extra_t_id_87>', '</extra_t_id_88>', '</extra_t_id_89>', '</extra_t_id_90>', '</extra_t_id_91>', '</extra_t_id_92>', '</extra_t_id_93>', '</extra_t_id_94>', '</extra_t_id_95>', '</extra_t_id_96>', '</extra_t_id_97>', '</extra_t_id_98>', '</extra_t_id_99>', '<loc_0>', '<loc_1>', '<loc_2>', '<loc_3>', '<loc_4>', '<loc_5>', '<loc_6>', '<loc_7>', '<loc_8>', '<loc_9>', '<loc_10>', '<loc_11>', '<loc_12>', '<loc_13>', '<loc_14>', '<loc_15>', '<loc_16>', '<loc_17>', '<loc_18>', '<loc_19>', '<loc_20>', '<loc_21>', '<loc_22>', '<loc_23>', '<loc_24>', '<loc_25>', '<loc_26>', '<loc_27>', '<loc_28>', '<loc_29>', '<loc_30>', '<loc_31>', '<loc_32>', '<loc_33>', '<loc_34>', '<loc_35>', '<loc_36>', '<loc_37>', '<loc_38>', '<loc_39>', '<loc_40>', '<loc_41>', '<loc_42>', '<loc_43>', '<loc_44>', '<loc_45>', '<loc_46>', '<loc_47>', '<loc_48>', '<loc_49>', '<loc_50>', '<loc_51>', '<loc_52>', '<loc_53>', '<loc_54>', '<loc_55>', '<loc_56>', '<loc_57>', '<loc_58>', '<loc_59>', '<loc_60>', '<loc_61>', '<loc_62>', '<loc_63>', '<loc_64>', '<loc_65>', '<loc_66>', '<loc_67>', '<loc_68>', '<loc_69>', '<loc_70>', '<loc_71>', '<loc_72>', '<loc_73>', '<loc_74>', '<loc_75>', '<loc_76>', '<loc_77>', '<loc_78>', '<loc_79>', '<loc_80>', '<loc_81>', '<loc_82>', '<loc_83>', '<loc_84>', '<loc_85>', '<loc_86>', '<loc_87>', '<loc_88>', '<loc_89>', '<loc_90>', '<loc_91>', '<loc_92>', '<loc_93>', '<loc_94>', '<loc_95>', '<loc_96>', '<loc_97>', '<loc_98>', '<loc_99>', '<loc_100>', '<loc_101>', '<loc_102>', '<loc_103>', '<loc_104>', '<loc_105>', '<loc_106>', '<loc_107>', '<loc_108>', '<loc_109>', '<loc_110>', '<loc_111>', '<loc_112>', '<loc_113>', '<loc_114>', '<loc_115>', '<loc_116>', '<loc_117>', '<loc_118>', '<loc_119>', '<loc_120>', '<loc_121>', '<loc_122>', '<loc_123>', '<loc_124>', '<loc_125>', '<loc_126>', '<loc_127>', '<loc_128>', '<loc_129>', '<loc_130>', '<loc_131>', '<loc_132>', '<loc_133>', '<loc_134>', '<loc_135>', '<loc_136>', '<loc_137>', '<loc_138>', '<loc_139>', '<loc_140>', '<loc_141>', '<loc_142>', '<loc_143>', '<loc_144>', '<loc_145>', '<loc_146>', '<loc_147>', '<loc_148>', '<loc_149>', '<loc_150>', '<loc_151>', '<loc_152>', '<loc_153>', '<loc_154>', '<loc_155>', '<loc_156>', '<loc_157>', '<loc_158>', '<loc_159>', '<loc_160>', '<loc_161>', '<loc_162>', '<loc_163>', '<loc_164>', '<loc_165>', '<loc_166>', '<loc_167>', '<loc_168>', '<loc_169>', '<loc_170>', '<loc_171>', '<loc_172>', '<loc_173>', '<loc_174>', '<loc_175>', '<loc_176>', '<loc_177>', '<loc_178>', '<loc_179>', '<loc_180>', '<loc_181>', '<loc_182>', '<loc_183>', '<loc_184>', '<loc_185>', '<loc_186>', '<loc_187>', '<loc_188>', '<loc_189>', '<loc_190>', '<loc_191>', '<loc_192>', '<loc_193>', '<loc_194>', '<loc_195>', '<loc_196>', '<loc_197>', '<loc_198>', '<loc_199>', '<loc_200>', '<loc_201>', '<loc_202>', '<loc_203>', '<loc_204>', '<loc_205>', '<loc_206>', '<loc_207>', '<loc_208>', '<loc_209>', '<loc_210>', '<loc_211>', '<loc_212>', '<loc_213>', '<loc_214>', '<loc_215>', '<loc_216>', '<loc_217>', '<loc_218>', '<loc_219>', '<loc_220>', '<loc_221>', '<loc_222>', '<loc_223>', '<loc_224>', '<loc_225>', '<loc_226>', '<loc_227>', '<loc_228>', '<loc_229>', '<loc_230>', '<loc_231>', '<loc_232>', '<loc_233>', '<loc_234>', '<loc_235>', '<loc_236>', '<loc_237>', '<loc_238>', '<loc_239>', '<loc_240>', '<loc_241>', '<loc_242>', '<loc_243>', '<loc_244>', '<loc_245>', '<loc_246>', '<loc_247>', '<loc_248>', '<loc_249>', '<loc_250>', '<loc_251>', '<loc_252>', '<loc_253>', '<loc_254>', '<loc_255>', '<loc_256>', '<loc_257>', '<loc_258>', '<loc_259>', '<loc_260>', '<loc_261>', '<loc_262>', '<loc_263>', '<loc_264>', '<loc_265>', '<loc_266>', '<loc_267>', '<loc_268>', '<loc_269>', '<loc_270>', '<loc_271>', '<loc_272>', '<loc_273>', '<loc_274>', '<loc_275>', '<loc_276>', '<loc_277>', '<loc_278>', '<loc_279>', '<loc_280>', '<loc_281>', '<loc_282>', '<loc_283>', '<loc_284>', '<loc_285>', '<loc_286>', '<loc_287>', '<loc_288>', '<loc_289>', '<loc_290>', '<loc_291>', '<loc_292>', '<loc_293>', '<loc_294>', '<loc_295>', '<loc_296>', '<loc_297>', '<loc_298>', '<loc_299>', '<loc_300>', '<loc_301>', '<loc_302>', '<loc_303>', '<loc_304>', '<loc_305>', '<loc_306>', '<loc_307>', '<loc_308>', '<loc_309>', '<loc_310>', '<loc_311>', '<loc_312>', '<loc_313>', '<loc_314>', '<loc_315>', '<loc_316>', '<loc_317>', '<loc_318>', '<loc_319>', '<loc_320>', '<loc_321>', '<loc_322>', '<loc_323>', '<loc_324>', '<loc_325>', '<loc_326>', '<loc_327>', '<loc_328>', '<loc_329>', '<loc_330>', '<loc_331>', '<loc_332>', '<loc_333>', '<loc_334>', '<loc_335>', '<loc_336>', '<loc_337>', '<loc_338>', '<loc_339>', '<loc_340>', '<loc_341>', '<loc_342>', '<loc_343>', '<loc_344>', '<loc_345>', '<loc_346>', '<loc_347>', '<loc_348>', '<loc_349>', '<loc_350>', '<loc_351>', '<loc_352>', '<loc_353>', '<loc_354>', '<loc_355>', '<loc_356>', '<loc_357>', '<loc_358>', '<loc_359>', '<loc_360>', '<loc_361>', '<loc_362>', '<loc_363>', '<loc_364>', '<loc_365>', '<loc_366>', '<loc_367>', '<loc_368>', '<loc_369>', '<loc_370>', '<loc_371>', '<loc_372>', '<loc_373>', '<loc_374>', '<loc_375>', '<loc_376>', '<loc_377>', '<loc_378>', '<loc_379>', '<loc_380>', '<loc_381>', '<loc_382>', '<loc_383>', '<loc_384>', '<loc_385>', '<loc_386>', '<loc_387>', '<loc_388>', '<loc_389>', '<loc_390>', '<loc_391>', '<loc_392>', '<loc_393>', '<loc_394>', '<loc_395>', '<loc_396>', '<loc_397>', '<loc_398>', '<loc_399>', '<loc_400>', '<loc_401>', '<loc_402>', '<loc_403>', '<loc_404>', '<loc_405>', '<loc_406>', '<loc_407>', '<loc_408>', '<loc_409>', '<loc_410>', '<loc_411>', '<loc_412>', '<loc_413>', '<loc_414>', '<loc_415>', '<loc_416>', '<loc_417>', '<loc_418>', '<loc_419>', '<loc_420>', '<loc_421>', '<loc_422>', '<loc_423>', '<loc_424>', '<loc_425>', '<loc_426>', '<loc_427>', '<loc_428>', '<loc_429>', '<loc_430>', '<loc_431>', '<loc_432>', '<loc_433>', '<loc_434>', '<loc_435>', '<loc_436>', '<loc_437>', '<loc_438>', '<loc_439>', '<loc_440>', '<loc_441>', '<loc_442>', '<loc_443>', '<loc_444>', '<loc_445>', '<loc_446>', '<loc_447>', '<loc_448>', '<loc_449>', '<loc_450>', '<loc_451>', '<loc_452>', '<loc_453>', '<loc_454>', '<loc_455>', '<loc_456>', '<loc_457>', '<loc_458>', '<loc_459>', '<loc_460>', '<loc_461>', '<loc_462>', '<loc_463>', '<loc_464>', '<loc_465>', '<loc_466>', '<loc_467>', '<loc_468>', '<loc_469>', '<loc_470>', '<loc_471>', '<loc_472>', '<loc_473>', '<loc_474>', '<loc_475>', '<loc_476>', '<loc_477>', '<loc_478>', '<loc_479>', '<loc_480>', '<loc_481>', '<loc_482>', '<loc_483>', '<loc_484>', '<loc_485>', '<loc_486>', '<loc_487>', '<loc_488>', '<loc_489>', '<loc_490>', '<loc_491>', '<loc_492>', '<loc_493>', '<loc_494>', '<loc_495>', '<loc_496>', '<loc_497>', '<loc_498>', '<loc_499>', '<loc_500>', '<other_0>', '<other_1>', '<other_2>', '<other_3>', '<other_4>', '<other_5>', '<other_6>', '<other_7>', '<other_8>', '<other_9>', '<other_10>', '<other_11>', '<other_12>', '<other_13>', '<other_14>', '<other_15>', '<other_16>', '<other_17>', '<other_18>', '<other_19>', '<other_20>', '<other_21>', '<other_22>', '<other_23>', '<other_24>', '<other_25>', '<other_26>', '<other_27>', '<other_28>', '<other_29>', '<other_30>', '<other_31>', '<other_32>', '<other_33>', '<other_34>', '<other_35>', '<other_36>', '<other_37>', '<other_38>', '<other_39>', '<other_40>', '<other_41>', '<other_42>', '<other_43>', '<other_44>', '<other_45>', '<other_46>', '<other_47>', '<other_48>', '<other_49>', '<other_50>', '<other_51>', '<other_52>', '<other_53>', '<other_54>', '<other_55>', '<other_56>', '<other_57>', '<other_58>', '<other_59>', '<other_60>', '<other_61>', '<other_62>', '<other_63>', '<other_64>', '<other_65>', '<other_66>', '<other_67>', '<other_68>', '<other_69>', '<other_70>', '<other_71>', '<other_72>', '<other_73>', '<other_74>', '<other_75>', '<other_76>', '<other_77>', '<other_78>', '<other_79>', '<other_80>', '<other_81>', '<other_82>', '<other_83>', '<other_84>', '<other_85>', '<other_86>', '<other_87>', '<other_88>', '<other_89>', '<other_90>', '<other_91>', '<other_92>', '<other_93>', '<other_94>', '<other_95>', '<other_96>', '<other_97>', '<other_98>', '<other_99>', '<other_100>', '<other_101>', '<other_102>', '<other_103>', '<other_104>', '<other_105>', '<other_106>', '<other_107>', '<other_108>', '<other_109>', '<other_110>', '<other_111>', '<other_112>', '<other_113>', '<other_114>', '<other_115>', '<other_116>', '<other_117>', '<other_118>', '<other_119>', '<other_120>', '<other_121>', '<other_122>', '<other_123>', '<other_124>', '<other_125>', '<other_126>', '<other_127>', '<other_128>', '<other_129>', '<other_130>', '<other_131>', '<other_132>', '<other_133>', '<other_134>', '<other_135>', '<other_136>', '<other_137>', '<other_138>', '<other_139>', '<other_140>', '<other_141>', '<other_142>', '<other_143>', '<other_144>', '<other_145>', '<other_146>', '<other_147>', '<other_148>', '<other_149>', '<other_150>', '<other_151>', '<other_152>', '<other_153>', '<other_154>', '<other_155>', '<other_156>', '<other_157>', '<other_158>', '<other_159>', '<other_160>', '<other_161>', '<other_162>', '<other_163>', '<other_164>', '<other_165>', '<other_166>', '<other_167>', '<other_168>', '<other_169>', '<other_170>', '<other_171>', '<other_172>', '<other_173>', '<other_174>', '<other_175>', '<other_176>', '<other_177>', '<other_178>', '<other_179>', '<other_180>', '<other_181>', '<other_182>', '<other_183>', '<other_184>', '<other_185>', '<other_186>', '<other_187>', '<other_188>', '<other_189>', '<other_190>', '<other_191>', '<other_192>', '<other_193>', '<other_194>', '<other_195>', '<other_196>', '<other_197>', '<other_198>', '<other_199>']\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "================              Failed loaded mae              Ignore this if you are finetuning the model              ================              \n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Some weights of UdopUnimodelForConditionalGeneration were not initialized from the model checkpoint at udop-unimodel-large-224 and are newly initialized: ['special_vis_token']\n",
      "You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "\n",
    "from core.models import UdopUnimodelForConditionalGeneration, UdopConfig, \\\n",
    "    UdopTokenizer\n",
    "\n",
    "# Change path to the model checkpoints\n",
    "config = UdopConfig.from_pretrained(\"udop-unimodel-large-224\")\n",
    "tokenizer = UdopTokenizer.from_pretrained(\"udop-unimodel-large-224\")\n",
    "model = UdopUnimodelForConditionalGeneration.from_pretrained(\"udop-unimodel-large-224\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "34f017fc",
   "metadata": {},
   "outputs": [],
   "source": [
    "from PIL import Image\n",
    "from core.common.utils import img_trans_torchvision, get_visual_bbox\n",
    "import json\n",
    "import numpy as np\n",
    "import torch\n",
    "\n",
    "# Load example document and process them\n",
    "im = Image.open('examples/00070353.png').convert('RGB')\n",
    "im = img_trans_torchvision(im).unsqueeze(0)\n",
    "visual_seg_data = get_visual_bbox(im.shape[-1]).unsqueeze(0)\n",
    "annot = json.load(open('examples/00070353.json'))\n",
    "\n",
    "all_bboxes = []\n",
    "all_text = []\n",
    "\n",
    "# Process task prefix, here we use document classification\n",
    "task_prefix = 'document classification.'\n",
    "sub_tokens = tokenizer.tokenize(task_prefix)\n",
    "for sub_token in sub_tokens:\n",
    "    all_text.append(sub_token)\n",
    "    all_bboxes.append([0, 0, 0, 0])\n",
    "        \n",
    "for item in annot['form']:\n",
    "    bb = np.array(item['box'])/1000.0\n",
    "    sub_tokens = tokenizer.tokenize(item['text'])\n",
    "    for sub_token in sub_tokens:\n",
    "        all_text.append(sub_token)\n",
    "        all_bboxes.append(bb)\n",
    "\n",
    "all_text_ids = tokenizer.convert_tokens_to_ids(all_text)\n",
    "input_ids = torch.tensor(np.array(all_text_ids)).unsqueeze(0)\n",
    "input_seg_data = torch.tensor(np.array(all_bboxes)).unsqueeze(0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "7bca352a",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/v-terrantang/anaconda3/envs/mmdif/lib/python3.8/site-packages/transformers/modeling_utils.py:812: FutureWarning: The `device` argument is deprecated and will be removed in v5 of Transformers.\n",
      "  warnings.warn(\n",
      "/home/v-terrantang/anaconda3/envs/mmdif/lib/python3.8/site-packages/transformers/modeling_utils.py:769: FutureWarning: The `device` argument is deprecated and will be removed in v5 of Transformers.\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "'form'"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Inference from the inputs\n",
    "output_ids = model.generate(\n",
    "            input_ids,\n",
    "            seg_data=input_seg_data,\n",
    "            image=im,\n",
    "            visual_seg_data=visual_seg_data,\n",
    "            use_cache=True,\n",
    "            decoder_start_token_id=None,\n",
    "            num_beams=1,\n",
    "            max_length=10,\n",
    "        )\n",
    "output_text = tokenizer.decode(output_ids[0][1:-1])\n",
    "output_text"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
